import functools

import albumentations as albu
import numpy as np


def to_tensor(x, **kwargs):
    return x.transpose(2, 0, 1).astype('float32')


def to_long(x, **kwargs):
    return x.astype("long")


def to_long_tensor(x, **kwargs):
    return x.transpose(2, 0, 1).astype("long")


def normalize(x, mean=None, std=None, **kwargs):
    if std is None:
        std = [0.229, 0.224, 0.225]
    if mean is None:
        mean = [0.485, 0.456, 0.406]
    x = x / 255.0
    mean = np.array(mean)
    x = x - mean
    std = np.array(std)
    x = x / std
    return x


def get_imagenet_params():
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    return {"mean": mean, "std": std}


def get_fn(pretrained="imagenet"):
    if pretrained == "imagenet":
        params = get_imagenet_params()
    return functools.partial(normalize, **params)


def get_preprocess(fn=get_fn()):
    _transform = [
        albu.Lambda(image=fn),
        albu.Lambda(image=to_tensor, mask=to_tensor),
    ]
    return albu.Compose(_transform)


def get_preprocess_2(fn=get_fn()):
    _transform = [
        albu.Lambda(image=fn),
        albu.Lambda(image=to_tensor, mask=to_long),
    ]
    return albu.Compose(_transform)


def base_aug():
    return albu.Compose([
        albu.HorizontalFlip(p=0.5),
        albu.VerticalFlip(p=0.5),
        albu.RandomRotate90(p=0.5),
    ])
